
/*
    TreeViewer - Cross-platform software to draw phylogenetic trees
    Copyright (C) 2023,2024  Giorgio Bianchini, University of Bristol
 
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published by
    the Free Software Foundation, version 3.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

using System;
using System.Collections.Generic;
using System.Linq;
using PhyloTree;
using TreeViewer;
using VectSharp;

namespace CircularCoordinates
{
    /// <summary>
    /// This module computes coordinates for the nodes of the tree in a "circular" style. The root node of the tree is placed at the
    /// center of the circle, and the tips are placed in rings whose diameter depends on the total distance from the root.
    /// </summary>
    /// 
    /// <description>
    /// <![CDATA[
    /// ## Further information
    /// 
    /// Here is an example of a tree drawn using circular coordinates (and with the appropriate shape for the _Branches_):
    /// 
    /// <p align="center">
    /// <img height="800" src="" />
    /// </p>
    /// ]]>
    /// </description>

    public static class MyModule
    {
        public const string Name = "Circular";
        public const string HelpText = "Computes the coordinates for a circular tree.";
        public const string Author = "Giorgio Bianchini";
        public static Version Version = new Version("1.2.0");
        public const string Id = "92aac276-3af7-4506-a263-7220e0df5797";
        public const ModuleTypes ModuleType = ModuleTypes.Coordinate;

        public static List<(string, string)> GetParameters(TreeNode tree)
        {
            double defaultRadius = tree.GetLeaves().Count * 20 / (2 * Math.PI);
            double innerRadius = double.IsNaN(tree.Length) ? defaultRadius * 0.1 : 0;

            defaultRadius = Math.Max(defaultRadius, 200);
            innerRadius = Math.Max(innerRadius, 20);

            System.Text.StringBuilder defaultSourceCode = new System.Text.StringBuilder();

            defaultSourceCode.AppendLine("using PhyloTree;");
            defaultSourceCode.AppendLine("using System.Collections.Generic;");
            defaultSourceCode.AppendLine("using TreeViewer;");
            defaultSourceCode.AppendLine("using VectSharp;");
            defaultSourceCode.AppendLine();
            defaultSourceCode.AppendLine("namespace a" + Guid.NewGuid().ToString("N"));
            defaultSourceCode.AppendLine("{");
            defaultSourceCode.AppendLine("\t//Do not change class name");
            defaultSourceCode.AppendLine("\tpublic static class CustomCoordinates");
            defaultSourceCode.AppendLine("\t{");
            defaultSourceCode.AppendLine("\t\t//Do not change method signature");
            defaultSourceCode.AppendLine("\t\tpublic static void GetCoordinates(TreeNode tree, ref Dictionary<string, Point> coordinates)");
            defaultSourceCode.AppendLine("\t\t{");
            defaultSourceCode.AppendLine("\t\t\t//TODO: change the coordinate values contained in the coordinates dictionary");
            defaultSourceCode.AppendLine("\t\t}");
            defaultSourceCode.AppendLine("\t}");
            defaultSourceCode.AppendLine("}");

            return new List<(string, string)>()
            {
                ( "Tree radius", "Group:2" ),
                
                /// <param name="Outer radius:">
                /// This parameter determines the radius of the ring that contains the tips with the highest distance from the root
                /// node (and, thus, the total diameter of the tree).
                /// </param>
                ( "Outer radius:", "NumericUpDown:" + defaultRadius.ToString(0) + "[\"0\",\"Infinity\"]" ),
                
                /// <param name="Inner radius:">
                /// This parameter determines the radius of the inner circle on which the root node is placed. If the radius is 0,
                /// the root node is placed exactly at the center of the tree; otherwise, it is placed on a ring with the specified
                /// radius.
                /// </param>
                ( "Inner radius:", "NumericUpDown:" + innerRadius.ToString(0) + "[\"0\",\"Infinity\"]" ),

                ( "Rotation", "Group:2" ),
                
                /// <param name="Rotation:">
                /// This parameter determines the orientation of the labels in the tree. Change this value to rotate all of the tips
                /// by the specified amount.
                /// </param>
                ( "Rotation:", "Slider:0[\"0\",\"360\",\"0°\"]" ),
                
                /// <param name="FixedRotations" display="Fixed rotations">
                /// These buttons can be used to set predefined values for the [Rotation](#rotation) of the tree.
                /// </param>
                ( "FixedRotations", "Buttons:[\"0°\",\"90°\",\"180°\",\"270°\"]" ),

                /// <param name="Sweep angle:">
                /// The angular width of the tree plot.
                /// </param>
                ( "Sweep angle:", "Slider:360[\"1\",\"360\",\"0°\"]" ),

                ( "Coordinate shift", "Group:5"),
                
                /// <param name="Coordinate shift:">
                /// This parameter determines the kind of coordinate shift that is applied. If the value is `None`, no coordinate
                /// shift is applied. If the value is `Relative`, the coordinates for each point are shifted by the amount specified
                /// by the selected [X](#x-attribute) and [Y](#y-attribute) attributes, relative to their default position. If the value is
                /// `Absolute`, the coordinates are set to the value specified by the selected [X](#x-attribute) and [Y](#y-attribute)
                /// attributes, regardless of their default position.
                /// </param>
                ( "Coordinate shift:", "ComboBox:0[\"None\",\"Relative\",\"Absolute\"]" ),

                /// <param name="X shift">
                /// If this check box is checked, the X coordinates of the tree nodes are shifted. Otherwise, they are left as is.
                /// </param>
                ( "X shift", "CheckBox:false"),
                
                /// <param name="X attribute:">
                /// This parameter determines the attribute used to shift the X coordinate of the points.
                /// </param>
                ( "X attribute:", "AttributeSelector:Length"),

                /// <param name="Y shift">
                /// If this check box is checked, the Y coordinates of the tree nodes are shifted. Otherwise, they are left as is.
                /// </param>
                ( "Y shift", "CheckBox:false"),
                
                /// <param name="Y attribute:">
                /// This parameter determines the attribute used to shift the Y coordinate of the points.
                /// </param>
                ( "Y attribute:", "AttributeSelector:Length"),

                /// <param name="Custom script:">
                /// This script can be used to modify the coordinate values.
                /// </param>
                ( "Custom script:", "SourceCode:" + defaultSourceCode.ToString()),
                
                /// <param name="Apply">
                /// This button applies the changes to the values of the other parameters and triggers a redraw of the tree.
                /// </param>
                ( "Apply", "Button:" )
            };
        }

        public static bool OnParameterChange(object tree, Dictionary<string, object> previousParameterValues, Dictionary<string, object> currentParameterValues, out Dictionary<string, ControlStatus> controlStatus, out Dictionary<string, object> parametersToChange)
        {
            controlStatus = new Dictionary<string, ControlStatus>()
            {

            };

            parametersToChange = new Dictionary<string, object>() { { "Apply", false }, { "FixedRotations", -1 } };

            int fixRot = (int)currentParameterValues["FixedRotations"];

            if ((int)currentParameterValues["Coordinate shift:"] == 0)
            {
                controlStatus["X shift"] = ControlStatus.Hidden;
                controlStatus["Y shift"] = ControlStatus.Hidden;
            }
            else
            {
                controlStatus["X shift"] = ControlStatus.Enabled;
                controlStatus["Y shift"] = ControlStatus.Enabled;
            }

            if ((int)currentParameterValues["Coordinate shift:"] != 0 && (bool)currentParameterValues["X shift"])
            {
                controlStatus["X attribute:"] = ControlStatus.Enabled;

                if (!(bool)previousParameterValues["X shift"] && (string)currentParameterValues["X attribute:"] == "Length")
                {
                    TreeNode treeNode = (TreeNode)tree;

                    string bestAttribute = "Length";

                    foreach (TreeNode leaf in treeNode.GetLeaves())
                    {
                        bool found = false;

                        foreach (KeyValuePair<string, object> kvp in leaf.Attributes)
                        {
                            if (kvp.Value is double)
                            {
                                if (kvp.Key.Equals("X", StringComparison.OrdinalIgnoreCase))
                                {
                                    bestAttribute = kvp.Key;
                                    found = true;
                                    break;
                                }
                                else if (kvp.Key.StartsWith("X", StringComparison.OrdinalIgnoreCase))
                                {
                                    bestAttribute = kvp.Key;
                                }
                            }
                        }

                        if (found)
                        {
                            break;
                        }
                    }

                    if (bestAttribute != "Length")
                    {
                        parametersToChange["X attribute:"] = bestAttribute;
                    }
                }
            }
            else
            {
                controlStatus["X attribute:"] = ControlStatus.Hidden;
            }

            if ((int)currentParameterValues["Coordinate shift:"] != 0 && (bool)currentParameterValues["Y shift"])
            {
                controlStatus["Y attribute:"] = ControlStatus.Enabled;

                if (!(bool)previousParameterValues["Y shift"] && (string)currentParameterValues["Y attribute:"] == "Length")
                {
                    TreeNode treeNode = (TreeNode)tree;

                    string bestAttribute = "Length";

                    foreach (TreeNode leaf in treeNode.GetLeaves())
                    {
                        bool found = false;

                        foreach (KeyValuePair<string, object> kvp in leaf.Attributes)
                        {
                            if (kvp.Value is double)
                            {
                                if (kvp.Key.Equals("Y", StringComparison.OrdinalIgnoreCase))
                                {
                                    bestAttribute = kvp.Key;
                                    found = true;
                                    break;
                                }
                                else if (kvp.Key.StartsWith("Y", StringComparison.OrdinalIgnoreCase))
                                {
                                    bestAttribute = kvp.Key;
                                }
                            }
                        }

                        if (found)
                        {
                            break;
                        }
                    }

                    if (bestAttribute != "Length")
                    {
                        parametersToChange["Y attribute:"] = bestAttribute;
                    }
                }
            }
            else
            {
                controlStatus["Y attribute:"] = ControlStatus.Hidden;
            }

            if (fixRot >= 0)
            {
                parametersToChange.Add("Rotation:", (double)(fixRot * 90));
                return true;
            }

            return (bool)currentParameterValues["Apply"] || currentParameterValues["Custom script:"] != previousParameterValues["Custom script:"];
        }

        public static Dictionary<string, Point> GetCoordinates(TreeNode tree, Dictionary<string, object> parameterValues)
        {
            int coordinateShift = (int)parameterValues["Coordinate shift:"];
            bool xShift = (bool)parameterValues["X shift"];
            bool yShift = (bool)parameterValues["Y shift"];
            string xShiftAttribute = (string)parameterValues["X attribute:"];
            string yShiftAttribute = (string)parameterValues["Y attribute:"];
            double sweepAngle = (double)parameterValues["Sweep angle:"] * Math.PI / 180;

            CompiledCode customScript = (CompiledCode)parameterValues["Custom script:"];

            List<TreeNode> nodes = tree.GetChildrenRecursive();

            Dictionary<string, Point> storedPos = new Dictionary<string, Point>();

            Dictionary<string, double> upstreamLengths = new Dictionary<string, double>();

            bool allNaN = true;

            foreach (TreeNode node in nodes)
            {
                if (node.Parent == null)
                {
                    upstreamLengths[node.Id] = double.IsNaN(node.Length) ? 0 : node.Length;
                }
                else
                {
                    if (!double.IsNaN(node.Length))
                    {
                        upstreamLengths[node.Id] = upstreamLengths[node.Parent.Id] + node.Length;
                        allNaN = false;
                    }
                    else
                    {
                        upstreamLengths[node.Id] = upstreamLengths[node.Parent.Id];
                    }
                }
            }

            if (allNaN)
            {
                foreach (TreeNode node in nodes)
                {
                    if (node.Parent == null)
                    {
                        upstreamLengths[node.Id] = double.IsNaN(node.Length) ? 0 : node.Length;
                    }
                    else
                    {
                        upstreamLengths[node.Id] = upstreamLengths[node.Parent.Id] + 1;
                    }
                }
            }

            double totalLength = upstreamLengths.Values.Max();

            Dictionary<string, double> yMultipliersCache = new Dictionary<string, double>();

            Dictionary<string, int> leafIndices = new Dictionary<string, int>();

            foreach (TreeNode leaf in tree.GetLeaves())
            {
                leafIndices[leaf.Id] = leafIndices.Count;
            }


            double getYMultiplier(TreeNode tree)
            {
                if (yMultipliersCache.TryGetValue(tree.Id, out double value))
                {
                    return value;
                }
                else
                {
                    if (tree.Children.Count == 0)
                    {
                        int ind = leafIndices[tree.Id];

                        double val = (ind + 0.5) / leafIndices.Count;
                        yMultipliersCache[tree.Id] = val;
                        return val;
                    }
                    else if (!tree.Attributes.ContainsKey("0c3400fd-8872-4395-83bc-a5dc5f4967fe") || (tree.Parent != null && tree.Parent.Attributes.ContainsKey("0c3400fd-8872-4395-83bc-a5dc5f4967fe")))
                    {
                        double val = 0.5 * (getYMultiplier(tree.Children[0]) + getYMultiplier(tree.Children[^1]));
                        yMultipliersCache[tree.Id] = val;
                        return val;
                    }
                    else
                    {
                        List<TreeNode> leaves = tree.GetLeaves();
                        double val = 0.5 * (getYMultiplier(leaves[0]) + getYMultiplier(leaves[^1]));
                        yMultipliersCache[tree.Id] = val;
                        return val;
                    }
                }
            }

            double outerRadius = (double)parameterValues["Outer radius:"];
            double innerRadius = (double)parameterValues["Inner radius:"];

            double rotation = (double)parameterValues["Rotation:"] * Math.PI / 180;

            for (int i = 0; i < nodes.Count; i++)
            {
                double r = upstreamLengths[nodes[i].Id] / totalLength * (outerRadius - innerRadius) + innerRadius;
                double theta = getYMultiplier(nodes[i]) * sweepAngle + rotation;

                double x = r * Math.Cos(theta);
                double y = r * Math.Sin(theta);

                storedPos[nodes[i].Id] = new Point(x, y);
            }

            if (!double.IsNaN(tree.Length))
            {
                double r = innerRadius;
                double theta = getYMultiplier(tree) * sweepAngle + rotation;

                double x = r * Math.Cos(theta);
                double y = r * Math.Sin(theta);

                storedPos[Modules.RootNodeId] = new Point(x, y);
            }
            else
            {
                storedPos[Modules.RootNodeId] = new Point(0, 0);
            }

            double scaleR = 1.0 / totalLength * (outerRadius - innerRadius);
            storedPos["d0ab64ba-3bcd-443f-9150-48f6e85e97f3"] = new Point(scaleR, innerRadius);

            storedPos[Id] = new Point(0, 0);

            if (coordinateShift > 0 && (xShift || yShift))
            {
                foreach (TreeNode node in tree.GetChildrenRecursiveLazy())
                {
                    Point pt = storedPos[node.Id];

                    double x = pt.X;
                    double y = pt.Y;

                    if (xShift)
                    {
                        if (node.Attributes.TryGetValue(xShiftAttribute, out object xAttribute) && xAttribute is double dX && !double.IsNaN(dX))
                        {
                            if (coordinateShift == 1)
                            {
                                x += dX;
                            }
                            else if (coordinateShift == 2)
                            {
                                x = dX;
                            }
                        }
                    }

                    if (yShift)
                    {
                        if (node.Attributes.TryGetValue(yShiftAttribute, out object yAttribute) && yAttribute is double dY && !double.IsNaN(dY))
                        {
                            if (coordinateShift == 1)
                            {
                                y += dY;
                            }
                            else if (coordinateShift == 2)
                            {
                                y = dY;
                            }
                        }
                    }

                    storedPos[node.Id] = new Point(x, y);
                }
            }

            object[] args = new object[] { tree, storedPos };

            Type customCoords = ModuleMetadata.GetTypeFromAssembly(customScript.CompiledAssembly, "CustomCoordinates");
            customCoords.InvokeMember("GetCoordinates", System.Reflection.BindingFlags.Default | System.Reflection.BindingFlags.InvokeMethod, null, null, args);

            storedPos = (Dictionary<string, Point>)args[1];

            return storedPos;
        }
    }
}

